package com.cfp4cloud.cfp.knowledge.support.rule;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.ReflectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.cfp4cloud.cfp.common.core.util.R;
import com.cfp4cloud.cfp.common.data.tenant.TenantContextHolder;
import com.cfp4cloud.cfp.knowledge.dto.AiMessageResultDTO;
import com.cfp4cloud.cfp.knowledge.dto.BaseAiRequest;
import com.cfp4cloud.cfp.knowledge.dto.ChatMessageDTO;
import com.cfp4cloud.cfp.knowledge.entity.AiChatRecordEntity;
import com.cfp4cloud.cfp.knowledge.mapper.AiChatRecordMapper;
import com.cfp4cloud.cfp.knowledge.service.AiAssistantService;
import com.cfp4cloud.cfp.knowledge.service.AiDataService;
import com.cfp4cloud.cfp.knowledge.support.constant.EmbedBizTypeEnums;
import com.cfp4cloud.cfp.knowledge.support.constant.ModelProviderFormatEnums;
import com.cfp4cloud.cfp.knowledge.support.feign.RemoteTableInfoService;
import com.cfp4cloud.cfp.knowledge.support.function.Chat2SqlFunctionCalling;
import com.cfp4cloud.cfp.knowledge.support.function.FunctionCalling;
import com.cfp4cloud.cfp.knowledge.support.provider.ChatMemoryAdvisorProvider;
import com.cfp4cloud.cfp.knowledge.support.provider.MemoryEmbeddingProvider;
import com.cfp4cloud.cfp.knowledge.support.provider.ModelProvider;
import com.cfp4cloud.cfp.knowledge.support.util.ChatMessageContextHolder;
import com.cfp4cloud.cfp.knowledge.support.util.JSONRepairUtil;
import com.cfp4cloud.cfp.knowledge.support.util.PromptBuilder;
import com.cfp4cloud.cfp.knowledge.support.util.ToolSpecificationsUtils;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Triple;
import org.jetbrains.annotations.Nullable;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import static com.cfp4cloud.cfp.knowledge.support.constant.AiPromptField.systemTime;
import static com.cfp4cloud.cfp.knowledge.support.provider.MemoryEmbeddingProvider.TEMP_ID;
import static dev.langchain4j.data.message.SystemMessage.systemMessage;
import static dev.langchain4j.data.message.UserMessage.userMessage;
import static dev.langchain4j.model.chat.request.ResponseFormatType.JSON;
import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

/**
 * 数据库聊天规则实现
 * <p>
 * 实现基于自然语言的数据库查询功能，将用户的自然语言转换为SQL语句并执行。 支持自动匹配数据集、生成图表、使用JSON Schema进行结构化输出等功能。
 *
 * @author chenda
 * @date 2024/3/26
 */
@Slf4j
@Component("databaseChat")
@RequiredArgsConstructor
public class DatabaseChatRule implements ChatRule {

	private final Chat2SqlFunctionCalling chat2SqlFunctionCalling;

	private final ChatMemoryAdvisorProvider chatMemoryAdvisorProvider;

	private final RemoteTableInfoService tableInfoService;

	private final AiChatRecordMapper chatRecordMapper;

	private final AiDataService aiDataService;

	private final ModelProvider modelProvider;

	/**
	 * 基于向量相似度自动匹配数据集
	 * <p>
	 * 当用户没有指定具体数据集时，通过向量搜索自动匹配语义最相似的数据集。 使用嵌入向量搜索CHAT2SQL类型的数据集。
	 * @param chatMessageDTO 聊天消息DTO，包含用户输入内容
	 * @return 如果找不到匹配的数据集返回错误提示，找到则返回null并更新上下文
	 */
	@Nullable
	private static Flux<AiMessageResultDTO> autoChoice(ChatMessageDTO chatMessageDTO) {
		EmbeddingSearchResult<TextSegment> searchResult = MemoryEmbeddingProvider.search(chatMessageDTO.getContent(),
				metadataKey(EmbedBizTypeEnums.Fields.type).isEqualTo(EmbedBizTypeEnums.CHAT2SQL.getType()));

		if (Objects.isNull(searchResult) || searchResult.matches().isEmpty()) {
			return Flux.just(new AiMessageResultDTO("未找到相关数据集建模，请点击下方+按钮选择目标数据集合"));
		}

		Long dataId = searchResult.matches().get(0).embedded().metadata().getLong(TEMP_ID);
		// 修改上下文中目标功能名称
		chatMessageDTO.getExtDetails().setDataId(dataId);
		ChatMessageContextHolder.set(chatMessageDTO);
		return null;
	}

	/**
	 * 构建聊天请求
	 * <p>
	 * 根据不同的模型提供商（OpenAI或其他）构建相应格式的聊天请求。 包含JSON Schema定义，确保模型输出符合预期的结构化格式。
	 * @param chatMemory 聊天记忆，存储历史对话
	 * @param functionCalling 函数调用接口，定义输出格式
	 * @param inputText 用户输入文本
	 * @param jsonModel JSON模型格式（OpenAI或其他）
	 * @return 构建好的聊天请求对象
	 */
	private ChatRequest buildChatRequest(ChatMemory chatMemory, FunctionCalling functionCalling, String inputText,
			String jsonModel) {
		ChatRequest.Builder builder = ChatRequest.builder();
		Class genericType = functionCalling.getGenericType();
		Field[] fields = ReflectUtil.getFields(genericType);
		JsonObjectSchema jsonObjectSchema = ToolSpecificationsUtils.parametersFrom(fields);
		JsonSchema jsonSchema = JsonSchema.builder()
			.name(functionCalling.functionName())
			.rootElement(jsonObjectSchema)
			.build();

		String metadata = PromptBuilder.render("knowledge-func-metadata.st", Map.of(BaseAiRequest.Fields.messageKey,
				ChatMessageContextHolder.get().getMessageKey(), systemTime, DateUtil.now()));

		// JSON Schema
		if (Objects.equals(jsonModel, ModelProviderFormatEnums.OPENAI.getFormat())) {
			chatMemory.add(userMessage(inputText + StrUtil.LF + metadata));
			List<ChatMessage> chatMessages = new ArrayList<>();
			chatMessages.add(systemMessage(PromptBuilder.render("ocr-system-json.st")));
			chatMessages.addAll(chatMemory.messages());
			builder.messages(chatMessages)
				.responseFormat(ResponseFormat.builder().type(JSON).jsonSchema(jsonSchema).build());
		}
		else {
			chatMemory.add(userMessage(PromptBuilder.render("user-json.st",
					Map.of("userInput", inputText + StrUtil.LF + metadata, "jsonSchema", jsonObjectSchema))));
			List<ChatMessage> chatMessages = new ArrayList<>();
			chatMessages.add(systemMessage(PromptBuilder.render("ocr-system-json.st")));
			chatMessages.addAll(chatMemory.messages());
			builder.messages(chatMessages);
		}

		return builder.build();
	}

	/**
	 * 处理数据库聊天请求
	 * <p>
	 * 主要流程： 1. 如果用户未指定数据集，自动匹配合适的数据集 2. 获取数据集的表结构信息 3. 使用AI模型将自然语言转换为SQL 4.
	 * 执行SQL并返回结果，支持图表展示
	 * @param chatMessageDTO 聊天上下文信息
	 * @return AI响应结果流，包含查询结果和可选的图表信息
	 */
	@Override
	public Flux<AiMessageResultDTO> process(ChatMessageDTO chatMessageDTO) {

		// 处理mcp，如果用户没有传递，则根据用户语义查询一个 data 数据集
		if (Objects.isNull(chatMessageDTO.getExtDetails())
				|| Objects.isNull(chatMessageDTO.getExtDetails().getDataId())) {
			Flux<AiMessageResultDTO> just = autoChoice(chatMessageDTO);
			if (just != null)
				return just;
		}

		Long dataId = chatMessageDTO.getExtDetails().getDataId();
		String tableSchemas = aiDataService.queryDataSchema(dataId);
		chatMessageDTO.getExtDetails().setDataId(dataId);
		chatMessageDTO.getExtDetails().setFuncName(chat2SqlFunctionCalling.functionName());
		ChatMessageContextHolder.set(chatMessageDTO);

		// 更新record 记录
		AiChatRecordEntity recordEntity = new AiChatRecordEntity();
		recordEntity.setRecordId(chatMessageDTO.getMessageKey());
		recordEntity.setExtDetails(JSONUtil.toJsonStr(chatMessageDTO.getExtDetails()));
		chatRecordMapper.updateById(recordEntity);

		Triple<ChatModel, AiAssistantService, String> jsonAssistantTriple = modelProvider
			.getAiJSONAssistant(chatMessageDTO.getModelName());

		ChatMemory chatMemory = chatMemoryAdvisorProvider.get(chatMessageDTO.getConversationId());

		String render = PromptBuilder.render("chat2db.st", Map.of("tableSchema", tableSchemas, "userInput",
				chatMessageDTO.getContent(), "tenantId",
				Objects.nonNull(TenantContextHolder.getTenantId()) ? TenantContextHolder.getTenantId() : StrUtil.EMPTY,
				systemTime, DateUtil.now()));

		ChatRequest chatRequest = buildChatRequest(chatMemory, chat2SqlFunctionCalling, render,
				jsonAssistantTriple.getRight());
		ChatResponse chatResponse = jsonAssistantTriple.getLeft().chat(chatRequest);

		String repair = JSONRepairUtil.repair(chatResponse.aiMessage().text());
		log.info("json chatResponse: {}", chatResponse.aiMessage().text());

		// 使用参数调用原有的函数逻辑
		ToolExecutionRequest toolExecutionRequest = ToolExecutionRequest.builder()
			.name(chat2SqlFunctionCalling.functionName())
			.arguments(repair)
			.id(chatMessageDTO.getConversationId())
			.build();

		R<String> resultR = chat2SqlFunctionCalling.execute(toolExecutionRequest);
		String result = StrUtil.isBlank(resultR.getData()) ? resultR.getMsg() : resultR.getData();
		chatMemory.add(AiMessage.from(chatResponse.aiMessage().text() + StrUtil.LF + result));

		ChatMessageDTO resultChatMessageDTO = ChatMessageContextHolder.get();

		AiMessageResultDTO aiMessageResultDTO = new AiMessageResultDTO(result);
		if (Objects.nonNull(resultChatMessageDTO) && Objects.nonNull(resultChatMessageDTO.getExtDetails())
				&& Objects.nonNull(resultChatMessageDTO.getExtDetails().getChartType())) {
			aiMessageResultDTO.setChartType(resultChatMessageDTO.getExtDetails().getChartType());
			aiMessageResultDTO.setChartId(resultChatMessageDTO.getExtDetails().getChartId());
		}
		return Flux.just(aiMessageResultDTO);
	}

}
